import cv2
import os
import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from keras.models import Sequential
from keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, BatchNormalization, Dropout
from keras.utils import np_utils
import tkinter as tk
from PIL import ImageTk, Image
from keras.preprocessing.image import ImageDataGenerator
from keras.applications.vgg16 import VGG16
from tensorflow.keras.applications.vgg16 import preprocess_input
from tensorflow.keras.models import Model


# 创建Tkinter应用程序窗口
window = tk.Tk()
window.title("昆虫分类器")
window.geometry("400x350")

# 标签和输入框
label = tk.Label(window, text="输入测试图片的文件名:")
label.pack()
input_entry = tk.Entry(window)
input_entry.pack()

# 显示分类结果的标签
result_label = tk.Label(window, text="")
result_label.pack()

# 加载数据集和标签
dataset_path = 'dataset/'
butterfly_prefix = 'fly'
beetle_prefix = 'jin'
snail_prefix = 'wo'
cockroach_prefix = 'zhang'
spider_prefix = 'zhizhu'

# 读取和预处理图像
def preprocess_image(image_path):
    image = cv2.imread(image_path)
    resized_image = cv2.resize(image, (224, 224))  # 调整图像大小为224x224像素
    normalized_image = resized_image / 255.0  # 归一化图像像素值
    return normalized_image

# 加载数据集和标签
def load_dataset():
    images = []
    labels = []
    for filename in os.listdir(dataset_path):
        if filename.startswith(butterfly_prefix):
            images.append(preprocess_image(os.path.join(dataset_path, filename)))
            labels.append('butterfly')
        elif filename.startswith(beetle_prefix):
            images.append(preprocess_image(os.path.join(dataset_path, filename)))
            labels.append('beetle')
        elif filename.startswith(snail_prefix):
            images.append(preprocess_image(os.path.join(dataset_path, filename)))
            labels.append('snail')
        elif filename.startswith(cockroach_prefix):
            images.append(preprocess_image(os.path.join(dataset_path, filename)))
            labels.append('cockroach')
        elif filename.startswith(spider_prefix):
            images.append(preprocess_image(os.path.join(dataset_path, filename)))
            labels.append('spider')

    # 转换为NumPy数组
    images = np.array(images)
    labels = np.array(labels)

    # 对标签进行编码
    label_encoder = LabelEncoder()
    labels = label_encoder.fit_transform(labels)

    return images, labels

# 训练分类器
def train_classifier(images, labels):
    # 调整图像大小以适应预训练的 VGG16 模型
    images_resized = np.array([cv2.resize(img, (224, 224)) for img in images])

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(images_resized, labels, test_size=0.2, random_state=42)

    # 将标签转换为分类独热编码
    num_classes = len(np.unique(labels))
    y_train_one_hot = np_utils.to_categorical(y_train, num_classes)
    y_test_one_hot = np_utils.to_categorical(y_test, num_classes)

    # 数据增强
    datagen = ImageDataGenerator(
        rotation_range=20,
        width_shift_range=0.2,
        height_shift_range=0.2,
        horizontal_flip=True
    )
    datagen.fit(X_train)

    # 使用预训练的 VGG16 模型进行特征提取
    base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
    X_train_features = base_model.predict(X_train)
    X_test_features = base_model.predict(X_test)

    # 将特征展平为一维向量
    X_train_flatten = X_train_features.reshape(X_train_features.shape[0], -1)
    X_test_flatten = X_test_features.reshape(X_test_features.shape[0], -1)

    # 训练 KNN 分类器
    knn_classifier = KNeighborsClassifier(n_neighbors=5)
    knn_classifier.fit(X_train_flatten, y_train)

    # 训练 SVM 分类器
    svm_classifier = SVC(kernel='linear')
    svm_classifier.fit(X_train_flatten, y_train)

    # 评估分类器的准确度
    knn_accuracy = knn_classifier.score(X_test_flatten, y_test)
    svm_accuracy = svm_classifier.score(X_test_flatten, y_test)

    print("KNN分类器准确度：", knn_accuracy)
    print("SVM分类器准确度：", svm_accuracy)

    # 构建更复杂的 CNN 模型，并使用相同的输入大小 (224, 224, 3)
    model = Sequential()
    model.add(Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)))
    model.add(BatchNormalization())
    model.add(MaxPooling2D((2, 2)))

    model.add(Conv2D(64, (3, 3), activation='relu'))
    model.add(BatchNormalization())
    model.add(MaxPooling2D((2, 2)))

    model.add(Conv2D(128, (3, 3), activation='relu'))
    model.add(BatchNormalization())
    model.add(MaxPooling2D((2, 2)))

    model.add(Flatten())
    model.add(Dense(128, activation='relu'))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes, activation='softmax'))

    # 编译模型，使用较小的学习率
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

    # 训练模型，使用数据增强
    model.fit(datagen.flow(X_train, y_train_one_hot, batch_size=32),
              epochs=50,
              validation_data=(X_test, y_test_one_hot))

    # 提取 CNN 模型的特征向量
    X_train_cnn_features = model.predict(X_train)
    X_test_cnn_features = model.predict(X_test)

    # 将独热编码的标签转换回原始标签形式
    y_train_labels = np.argmax(y_train_one_hot, axis=1)
    y_test_labels = np.argmax(y_test_one_hot, axis=1)

    # 训练基于 CNN 特征的 KNN 分类器
    knn_classifier_cnn = KNeighborsClassifier(n_neighbors=5)
    knn_classifier_cnn.fit(X_train_cnn_features, y_train_labels)

    # 训练基于 CNN 特征的 SVM 分类器
    svm_classifier_cnn = SVC(kernel='linear')
    svm_classifier_cnn.fit(X_train_cnn_features, y_train_labels)

    return model, knn_classifier, svm_classifier, knn_classifier_cnn, svm_classifier_cnn

# 预测昆虫类型
# 加载预训练的VGG16模型，不包括顶层全连接层
base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
# 添加一个全局平均池化层来替代全连接层
# model = Model(inputs=base_model.input, outputs=base_model.output)
def preprocess_image(image_path):
    image = cv2.imread(image_path)
    image = cv2.resize(image, (224, 224))  # 调整图像大小为224x224
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = np.array(image, dtype=np.float32)
    image = preprocess_input(image)  # 使用VGG16的预处理函数
    return image
def predict_insect(image_path, model, knn_classifier, svm_classifier):
    test_image = preprocess_image(image_path)
    test_image = np.expand_dims(test_image, axis=0)  # 扩展维度以适应模型输入

    # 使用CNN模型进行预测
    cnn_prediction = model.predict(test_image)
    cnn_predicted_class = np.argmax(cnn_prediction)

    # 使用VGG16模型提取特征
    # feature_extractor = model.layers[-2]  # 假设使用倒数第二层作为特征提取层
    # features = feature_extractor(test_image)
    # features = features.numpy()  # 转换为numpy数组
    features = base_model.predict(test_image)
    features_flatten = features.reshape(features.shape[0], -1)  # 展平特征

    # 使用KNN进行预测
    knn_prediction = knn_classifier.predict(features_flatten)

    # 使用SVM进行预测
    svm_prediction = svm_classifier.predict(features_flatten)

    return cnn_predicted_class, knn_prediction[0], svm_prediction[0]

# 加载数据集和标签
images, labels = load_dataset()

# 训练分类器
model, knn_classifier, svm_classifier, knn_classifier_cnn, svm_classifier_cnn = train_classifier(images, labels)

# 预测按钮的回调函数
def predict():
    # 获取输入的图片文件名
    image_filename = input_entry.get()
    image_path = os.path.join(dataset_path, image_filename)

    if os.path.exists(image_path):
        # 预测测试图像
        cnn_predicted_class, knn_prediction, svm_prediction = predict_insect(image_path, model, knn_classifier, svm_classifier)

        # 显示预测结果
        class_mapping = {0: 'butterfly', 1: 'beetle', 2: 'snail', 3: 'cockroach', 4: 'spider'}
        cnn_predicted_insect = class_mapping[cnn_predicted_class]
        knn_predicted_insect = class_mapping[knn_prediction]
        svm_predicted_insect = class_mapping[svm_prediction]
        result_label.configure(text=f'CNN分类器预测结果: {cnn_predicted_insect}\n'
                                    f'KNN分类器预测结果: {knn_predicted_insect}\n'
                                    f'SVM分类器预测结果: {svm_predicted_insect}')

        # 显示测试图片
        img = Image.open(image_path)
        img = img.resize((200, 200), Image.LANCZOS)
        img = ImageTk.PhotoImage(img)
        image_label.configure(image=img)
        image_label.image = img  # 保持对图片对象的引用，避免被垃圾回收

    else:
        result_label.configure(text="文件路径不存在！")


# 预测按钮
predict_button = tk.Button(window, text="预测", command=predict)
predict_button.pack()

# 显示测试图片的标签
image_label = tk.Label(window)
image_label.pack()

# 运行Tkinter应用程序的主循环
window.mainloop()
